# -*- coding:utf-8 -*-
from pathlib import Path
import frida
import lldb
import re
import shlex
import ast
import paramiko
import _thread
import os
import time

path = 'lldb_info'
frida_manager = None
proxy = 1111
port = 2222
hostname = '127.0.0.1'
username = 'root'
password = 'alpine'

# [+] 正确提示
# [-] 错误提示

js_code = """

rpc.exports =
{
    path: function ()
    {
       get_app_path()
    }
};

function get_app_path()
{
    var app_path = ObjC.classes.NSBundle.mainBundle().bundlePath();
    console.log('[+] ' + app_path);
};

anti()
function anti()
{
    var module = Process.enumerateModules()[0];
    var ptrace = Module.getExportByName(module.name, 'ptrace');
    var originPtrace = new NativeFunction(ptrace, 'void', ["int", "pointer", "pointer", "int"]);
    Interceptor.replace(ptrace, new NativeCallback((arg1, arg2, arg3, arg4) =>
    {
       if (arg1 == 31)
       {
           console.log("[+] hook ptrace.");
           return 0;
       }
       else
       {
           return originPtrace(arg1, arg2, arg3, arg4);
       }
    }, 'void', ["int", "pointer", "pointer", "int"]));


   var dlsym = Module.getExportByName(module.name, 'dlsym')
   Interceptor.attach(dlsym,
   {
       onEnter(args)
       {
           this.path = false;
           var pointer = new NativePointer(args[1]);
           var arg = pointer.readUtf8String();
           if (arg.indexOf("ptrace") > -1)
           {
               console.log("[+] hook ptrace.");
               this.path = true;
           }
       },
       onLeave(retval)
       {
           if (this.path)
           {
               try
               {
                   var ptracePtr = ptr(retval);
                   var originPtrace = new NativeFunction(ptracePtr, 'void', ["int", "pointer", "pointer", "int"]);
                   Interceptor.replace(ptracePtr, new NativeCallback((arg1, arg2, arg3, arg4) =>
                   {
                    if (arg1 == 31)
                    {
                        console.log("[+] hook dlsym.");
                        return 0;
                    }
                    else
                    {
                        return originPtrace(arg1, arg2, arg3, arg4);
                    }
                    }, 'void', ["int", "pointer", "pointer", "int"]));
               }
               catch (error)
               {
                   console.log("[-] " + error);
               }
           }
       }
   });

    var syscall = Module.getExportByName(module.name, 'syscall');
    var originSyscall = new NativeFunction(syscall, 'int', ["int", "int"]);
    Interceptor.replace(syscall, new NativeCallback((arg1, arg2, arg3, arg4) =>
    {
        if (arg1 == 26 && arg2 == 31)
        {
            console.log("[+] hook syscall.");
            return 0;
        }
        else
        {
            return originSyscall(arg1, arg2);
        }
        return 0;
    }, 'int', ['int', 'int', 'int', 'int']));

    var sysctl = Module.getExportByName(module.name, 'sysctl');
    var OriginSysctl = new NativeFunction(sysctl, 'int', ["pointer", "int", "pointer", "pointer", "pointer", "int"]);
    Interceptor.replace(sysctl, new NativeCallback((name, name_len, info, info_len, newp, newp_len)=>
    {
        var ret = OriginSysctl(name, name_len, info, info_len, newp, newp_len);
        if(name_len == 4)
        {
            var p_name = new NativePointer(name);
            var name0 = p_name.readInt();
            var name1 = p_name.add(4).readInt();
            var name2 = p_name.add(8).readInt();

            if(name0 == 1 && name1 == 14 && name2 == 1)
            {
                console.log("[+] hook sysctl.");
                var p_info_size = new NativePointer(info_len)
                var info_size = p_info_size.readInt();
                // console.log(info_size);
                var p_info = new NativePointer(info)
                p_info.add(32).writeByteArray([0x00,0x00]);
                // console.log(p_info.readByteArray(info_size));
            }
        }
        return ret;
    }, 'int', ["pointer", "int", "pointer", "pointer", "pointer", "int"]));

    const pattern = '01 10 00 D4';
    try
    {
        Memory.scan(module.base, module.size, pattern,
        {
            onMatch(address, size)
            {
                try
                {
                    Interceptor.attach(address.sub(4),
                    {
                        onEnter: function (args)
                        {
                            Memory.protect(address, 4, 'rw-');
                            address.writeByteArray([0x1F, 0x20, 0x03, 0xD5]);
                        },
                        onLeave: function (retval)
                        {
                        }
                    });
                }
                catch(error)
                {
                    console.log(error);
                }
            },
            onComplete() {
            }
        });
    }
    catch (error)
    {
        console.log(error)
    }
}
"""


def on_message(message, data):
    if message['type'] == 'send':
        text = message['payload']
        print(text)
    else:
        print(message)


class IosDevicesManager:
    def __init__(self):
        self.ssh = paramiko.SSHClient()

    def connect(self):
        self.ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        try:
            self.ssh.connect(hostname=hostname, username=username, password=password, allow_agent=True,timeout=5, port=port)
        except Exception as e:
            print(e)

    def thread_(self, cmd, temp):
        try:
            stdin, stdout, stderr = self.ssh.exec_command(cmd)
            exec_result = stdout.readlines()
            for item in exec_result:
                print('[+] ' + item.replace('\n', ''))
        except Exception as e:
            print(e)

    def exec_asy_cmd(self, cmd):
        try:
            _thread.start_new_thread(self.thread_, (cmd, 1))
        except Exception as e:
            print(e)

    def exec_syn_cmd(self, cmd):
        try:
            stdin, stdout, stderr = self.ssh.exec_command(cmd)
            exec_result = stdout.readlines()
            for item in exec_result:
                print('[+] ' + item.replace('\n', ''))
        except Exception as e:
            print(e)


class FridaManager:
    def __init__(self):
        self._script = None
        self._device = None
        self._session = None
        self.init_device()

    def init_device(self):
        try:
            self._device = frida.get_usb_device()
        except Exception as e:
            self._device = frida.get_remote_device()

    def attach(self, text):
        self._session = self._device.attach(text)
        self._script = self._session.create_script(js_code)
        self._script.on('message', on_message)
        self._script.load()

    def swap(self, text):
        pid = self._device.spawn(text)
        self._session = self._device.attach(pid)
        self._script = self._session.create_script(js_code)
        self._script.on('message', on_message)
        self._script.load()
        self._device.resume(pid)
        time.sleep(2)

    def detach(self):
        if self._session:
            self._session.detach()

    def get_exports(self):
        return self._script.exports

    def get_applications(self):
        try:
            if not self._device:
                return None
            applications = self._device.enumerate_applications()
            if len(applications) > 0:
                for application in sorted(applications, key=self.cmp_to_key(self.compare_applications)):
                    if application.pid != 0:
                        print('[+]', application.pid, application.name, application.identifier)
                    else:
                        print('[+]', '-', application.name, application.identifier)
        except Exception as e:
            print(e)

    @staticmethod
    def compare_applications(a, b):
        a_is_running = a.pid != 0
        b_is_running = b.pid != 0
        if a_is_running == b_is_running:
            if a.name > b.name:
                return 1
            elif a.name < b.name:
                return -1
            else:
                return 0
        elif a_is_running:
            return -1
        else:
            return 1

    @staticmethod
    def cmp_to_key(cmp):
        """Convert a cmp= function into a key= function"""

        class K:
            def __init__(self, obj):
                self.obj = obj

            def __lt__(self, other):
                return cmp(self.obj, other.obj) < 0

            def __gt__(self, other):
                return cmp(self.obj, other.obj) > 0

            def __eq__(self, other):
                return cmp(self.obj, other.obj) == 0

            def __le__(self, other):
                return cmp(self.obj, other.obj) <= 0

            def __ge__(self, other):
                return cmp(self.obj, other.obj) >= 0

            def __ne__(self, other):
                return cmp(self.obj, other.obj) != 0

        return K


def command_exc(cmd):
    """
    静默执行lldb命令不展示
    :param cmd:
    :return:
    """
    interpreter = lldb.debugger.GetCommandInterpreter()
    return_object = lldb.SBCommandReturnObject()
    interpreter.HandleCommand(cmd, return_object)  # 执行命令
    output = return_object.GetOutput()  # 获取命令的返回值
    return output


def is_hex_number(address_info):
    """
    判断输入地址信息是否为0x开头
    :param address_info: 地址
    :return:
    """
    match = re.match(r'.+(0x[0-9a-fA-F]+)', address_info)  # 正则匹配(0x开头)
    if match:
        result = match.group(1)
    else:
        result = '0x' + match.group(1)
    return result


def get_aslr():
    """
    获取随机地址
    """
    output = command_exc('image list -o')
    match = re.match(r'.+(0x[0-9a-fA-F]+)', output)
    if match:
        return match.group(1)
    else:
        return None


def get_base_address():
    """
     获取模块基地址
     """
    output = command_exc('image list -h')
    match = re.match(r'.+(0x[0-9a-fA-F]+)', output)
    if match:
        return match.group(1)
    else:
        return None


def get_commands(command):
    commands = []
    if command:
        temps = command.split(' ')
        for item in temps:
            if item:
                commands.append(item)
    return commands


def breakpoint_save(bps):
    """
     保存断点
    :param bps: 断点列表
    """
    file_lldb = open(path, 'w')
    file_lldb.write(str(bps))
    file_lldb.write(' ')
    file_lldb.close()


def breakpoint_read():
    """
    读取断点
    """
    try:
        file_lldb = open(path, 'r')
        data = ast.literal_eval(file_lldb.read())
        file_lldb.close()
    except IOError as e:
        print("fils is not found!")
        return {}
    return data


def breakpoint_write(bp):
    """
    写入断点
    """
    key = list(bp.keys())[0]
    value = list(bp.values())[0]
    file_lldb = dict(breakpoint_read())
    file_lldb[key] = value
    breakpoint_save(file_lldb)


def get_lldb_bps():
    """
    获取lldb断点列表
    """
    output = command_exc('breakpoint  list -b')
    it = re.finditer(r'.+(0x[0-9a-fA-F]+)', output)
    bps = {}
    for item in it:
        bps[item.group(0).split(':')[0]] = item.group(1)
    return bps


def port_forward(portx, porty):
    exc_cmd('iproxy ' + str(portx) + ' ' + str(porty))


def exc_cmd(cmd):
    result = os.popen(cmd)
    res = result.read()
    result = []
    for line in res.splitlines():
        line.replace('\n', '')
    return result


def offset(debugger, command, result, internal_dict):
    try:
        cmd = get_commands(command)
        base_address = get_base_address()
        if not cmd:
            print('[+]', base_address)
            return
        base = int(base_address, 16)
        address = int(cmd[0], 16)

        if address > base:
            print('[+]', hex(address-base))
        else:
            print('[-]', 'address input is error.')
    except Exception as e:
        print('[-]', 'param not is number.')


def aslr(debugger, command, result, internal_dict):
    try:
        cmd = get_commands(command)
        aslr_address = get_aslr()
        if not cmd:
            print('[+]', aslr_address)
            return
        aslr = int(aslr_address, 16)
        address = int(cmd[0], 16)
        print('[+]', hex(address+aslr))
    except Exception as e:
        print('[-]', 'param not is number.')


def breakpoint_from_lldb_read(debugger, command, result, internal_dict):
    """
    读取lldb断点
    """
    bps = get_lldb_bps()
    for item in bps:
        print(item, bps.get(item))


def breakpoint_from_file_read(debugger, command, result, internal_dict):
    """
    读取文件断点
    """
    bps = breakpoint_read()
    for item in bps:
        print(item, bps.get(item))


def breakpoint_clear(debugger, command, result, internal_dict):
    """
    清空断点
    """
    bps = breakpoint_read()
    for item in bps:
        debugger.HandleCommand('breakpoint delete "%s"' % item)

    file_lldb = open(path, 'w')
    file_lldb.close()


def breakpoint_reset(debugger, command, result, internal_dict):
    """
    重新设置断点
    """
    bps = breakpoint_read()
    aslr = get_aslr()
    for item in bps:
        debugger.HandleCommand('breakpoint set -a "%s+%s"' % (aslr, bps.get(item)))
    else:
        print('[-]', 'ASLR not found or file offset is null!')


def breakpoint_delete(debugger, command, result, internal_dict):
    """
    删除断点
    """
    bpid = shlex.split(command)[0]
    if not bpid:
        print('[-]', 'Please input the breakpoint id!')
        return

    bps = breakpoint_read()
    for item in bps:
        if bpid == item:
            debugger.HandleCommand('breakpoint delete "%s"' % (bpid))
            del bps[bpid]
            breakpoint_save(bps)
            break


def breakpoint_address(debugger, command, result, internal_dict):
    """
    打断点
    """
    offset = shlex.split(command)[0]
    if not offset:
        print('[-]', 'Please input the address!')
        return

    bps = breakpoint_read()
    offsets = bps.values()
    for item in offsets:
        if item == offset:
            print('[-]', 'Breakpoint repeat!')
            return

    aslr = get_aslr()
    if aslr:
        debugger.HandleCommand('breakpoint set -a "%s+%s"' % (aslr, offset))
        breakpoint_save(get_lldb_bps())
    else:
        print('[-]','ASLR not found!')


def breakpoint_list(debugger, command, result, internal_dict):
    """
    查看断点列表
    """
    debugger.HandleCommand('breakpoint list')


def debug_p(debugger, command, result, internal_dict):
    """
    单步步过
    """
    debugger.HandleCommand('ni')


def debug_t(debugger, command, result, internal_dict):
    """
    单步步入
    """
    debugger.HandleCommand('si')


def debug_g(debugger, command, result, internal_dict):
    """
    运行
    """
    debugger.HandleCommand('c')


def attach(debugger, command, result, internal_dict):
    """
    附加连接
    """
    debugger.HandleCommand('process connect connect://127.0.0.1:%s' % proxy)


def proxy_forward(debugger, command, result, internal_dict):
    exc_cmd('killall iproxy')
    _thread.start_new_thread(port_forward, (port, 22))
    _thread.start_new_thread(port_forward, (proxy, proxy))


def app(debugger, command, result, internal_dict):
    """
    get app
    """
    frida_manager.get_applications()


def debug(debugger, command, result, internal_dict):
    """
    debugserver 附加app
    """
    manager.exec_syn_cmd('killall debugserver')

    cmds = get_commands(command)
    if not cmds:
        print('[-] cmd is null.')
        return

    cmd = cmds[0]
    # /var/containers/Bundle/Application/A6ECEC4E-9244-49F0-9D2A-82AB2AAFF176/hxoffice-ios.app
    # /var/containers/Bundle/Application/A47B6F3F-CC9D-4E3F-A7AC-4F511440174C/WeChat.app
    if 'p' == cmd:
        if len(cmds) >= 2:
            pid = cmds[1]
            manager.exec_asy_cmd('debugserver 127.0.0.1:' + str(proxy) + ' -a ' + pid)
        else:
            print('[-] pid is null.')
    elif 'x' == cmd:
        if len(cmds) >= 2:
            path = cmds[1]
            manager.exec_asy_cmd('debugserver -x backboard 127.0.0.1:' + str(proxy) + ' ' + path)
        else:
            print('[-] app path is null.')
    else:
        print('[-] command error.')


def frida_ios(debugger, command, result, internal_dict):
    cmds = get_commands(command)
    if not cmds:
        print('[-] cmd is null.')
        return

    try:
        cmd = cmds[0]
        if 'path' == cmd or 'p' == cmd:
            frida_manager.get_exports().path()
        elif 'attach' == cmd or 'a' == cmd:
            if len(cmds) >= 3:
                if 'name' == cmds[1] or 'n' == cmds[1]:
                    frida_manager.attach(cmds[2])
                elif 'pid' == cmds[1] or 'p' == cmds[1]:
                    frida_manager.attach(int(cmds[2]))
                else:
                    print('')
        elif 'swap' == cmd or 's' == cmd:
            if len(cmds) >= 2:
                frida_manager.swap(cmds[1])
        elif 'detach' == cmd or 'd' == cmd:
            frida_manager.detach()
        else:
            print('[-] command error.')
    except Exception as e:
        print(e)


def ios_cmd(debugger, command, result, internal_dict):
    """
    执行cmd命令
    """
    while True:
        cmd = input("[cmd] ")
        if cmd != 'exit':
            manager.exec_syn_cmd(cmd)
        else:
            break


def init():
    global manager
    manager = IosDevicesManager()
    manager.connect()

    global frida_manager
    frida_manager = FridaManager()

# command script import ~/PycharmProjects/test/lldb_script.py


def __lldb_init_module(debugger, internal_dict):
    init()

    debugger.HandleCommand('command script add -f lldb_script.aslr aslr')
    debugger.HandleCommand('command script add -f lldb_script.offset offset')

    debugger.HandleCommand('command script add -f lldb_script.attach attach')

    debugger.HandleCommand('command script add -f lldb_script.breakpoint_address bp')
    debugger.HandleCommand('command script add -f lldb_script.breakpoint_list bl')
    debugger.HandleCommand('command script add -f lldb_script.breakpoint_delete bd')

    debugger.HandleCommand('command script add -f lldb_script.breakpoint_clear bclear')
    debugger.HandleCommand('command script add -f lldb_script.breakpoint_reset breset')
    debugger.HandleCommand('command script add -f lldb_script.breakpoint_from_file_read brf')
    debugger.HandleCommand('command script add -f lldb_script.breakpoint_from_lldb_read brl')

    debugger.HandleCommand('command script add -f lldb_script.debug_p p')
    debugger.HandleCommand('command script add -f lldb_script.debug_t t')
    debugger.HandleCommand('command script add -f lldb_script.debug_g g')

    debugger.HandleCommand('command script add -f lldb_script.app app')
    debugger.HandleCommand('command script add -f lldb_script.debug debug')
    debugger.HandleCommand('command script add -f lldb_script.ios_cmd cmd')
    debugger.HandleCommand('command script add -f lldb_script.frida_ios frida')
    # debugger.HandleCommand('command script add -f lldb_script.proxy_forward forward')


